package com.mrx.loader;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.Arrays;
import java.util.List;

/**
 * 自定义的 {@link ClassLoader}, 用以完成修改加载的类的字节码, 操纵类加载行为等功能
 */
public class AgentClassLoader extends URLClassLoader {

    private static final Logger logger = LoggerFactory.getLogger(AgentClassLoader.class);

    public AgentClassLoader() {
        super(new URL[]{});
    }

    /**
     * 将一批文件添加到 classPath
     *
     * @param libFiles 要添加的文件
     */
    @SuppressWarnings("unused")
    public void addFiles(List<File> libFiles) {
        try {
            for (File libFile : libFiles) {
                addURL(libFile.toURI().toURL());
            }
        } catch (MalformedURLException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 将一批 URL 添加到 classPath
     *
     * @param urls 要添加的 url
     */
    public void addURLs(List<URL> urls) {
        logger.debug("add classPath: {}", urls);
        List<URL> list = Arrays.asList(getURLs());
        urls.removeIf(list::contains);
        for (URL url : urls) {
            addURL(url);
        }
    }

    /**
     * 将一个 url 添加到 classPath
     *
     * @param url 要添加的 url
     */
    @Override
    public void addURL(URL url) {
        if (!Arrays.asList(getURLs()).contains(url)) {
            super.addURL(url);
        }
    }

    /**
     * 加载一个 class, 该行为与 {@link URLClassLoader#loadClass(String)} 一致, 这里重写只是为了在 stackTrace 里显示出来, 方便调试
     */
    @Override
    public Class<?> loadClass(String name) throws ClassNotFoundException {
        return super.loadClass(name);
    }

    /**
     * 使用字节码定义一个 class
     *
     * @param name       class 的名字
     * @param clazzBytes class 字节码
     * @return 定义好的 class
     */
    public Class<?> loadClass(String name, byte[] clazzBytes) {
        byte[] transformed = transform(name, clazzBytes);
        Class<?> clazz = defineClass(name, transformed, 0, transformed.length);
        logger.debug("加载 class: {}", clazz);
        return clazz;
    }

    /**
     * 使用 ASM 修改字节码, 替换部分字符串
     *
     * @param className       类名
     * @param classFileBuffer 类的字节码
     * @return 修改后的类的字节码
     */
    private byte[] transform(String className, byte[] classFileBuffer) {
        ClassReader cr = new ClassReader(classFileBuffer);
        ClassWriter cw = new ClassWriter(cr, 0);
        cr.accept(new TranslateClassVisitor(className, Opcodes.ASM7, cw), 0);
        return cw.toByteArray();
    }

}
